/**
 * Adaptive array definition
 *
 * \file adaptive-array.hpp
 * \author tomby
 * \date 04-01-2009
 **/

/*  Copyright (C) 2009  Tomas 'tomby' Bily

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>. */

#ifndef __DBC_ADAPTIVE_ARRAY__
#define __DBC_ADAPTIVE_ARRAY__

namespace dbc {

  /**
   * Adaptive array class
   * 
   *  Adaptive array is multi-dimensional two layer array. First layer
   *  is static array of nodes (struct Node). Nodes contain dynamicaly
   *  refinable arrays (DRA) of values (blitz::Array). DRA is
   *  multi-dimensional array with sizes of power of two in each
   *  dimension (Hypercube). Refinemment of each DRA is bounded by
   *  constat.
   *
   * Example 1D:
   *
   *  |Node 1|Node 2|Node 3|Node 4|
   *     |      |      |      |
   *     #(1)  ####(4) ##(2) ########(8)
   *
   */

  template <typename ref_coars_constr>
  class AdaptiveArray
  {
  public:
    typedef ref_coars_constr AdaptiveConstrain;
    typedef typename ref_coars_constr::NumericType	NumericType;

    enum		{ Dim = ref_coars_constr::Dim };

    typedef blitz::Array<NumericType, Dim> NumericArrayType;
    typedef blitz::TinyVector<int, Dim> IndexType;

  public:

    /**
     * Structure node 
     */
    struct Node
    {
      /// Level of refinement. (0 - coarsest, RefineLevel - finest)
      int		level;
      /// Data array with size of 1 << level in each dimension
      NumericArrayType	data;

    public:
      Node ()
	: level (0)
      {}

      Node (int ilevel, NumericArrayType& idata)
	: level (ilevel)
      { data.reference (idata); }

      Node (int ilevel)
	: level (ilevel), data (1 << ilevel)
      {}

      ~Node ()
      {
	/***/;
      }
    };

  public:
    typedef Node NodeType;
    typedef blitz::Array<Node*, Dim> ArrayType;
    
  private:
    int RefineLevel;
    NumericArrayType data_array;
    ArrayType node_array;
    
    IndexType size;

  public:
    class cell_index
    {
      IndexType node_idx, val_idx;

    public:
      cell_index () {}
      cell_index (IndexType node, IndexType val);

      IndexType& node ();
      IndexType& val ();
    };

  public:
    class node_iterator;

    class level_iterator 
    {
      RectDomIterator<ref_coars_constr::Dim> pos;
      node_iterator it;
      int level_up;
	
    public:
      level_iterator (node_iterator& iit, int ilevel_up);
	
      level_iterator& operator++ ();
      NumericArrayType operator* ();
      
      bool end_p ();
    };

  public:
    class node_iterator
    {
      IndexType		pos;
      AdaptiveArray*	aa;

    public:
      node_iterator (AdaptiveArray* iaa)
	: pos (0), aa (iaa)
      {}

      node_iterator (AdaptiveArray* iaa, bool)
	: pos (iaa->dims () - 1), aa (iaa)
      {}

      node_iterator (AdaptiveArray* iaa, const IndexType& ipos)
	: pos (ipos), aa (iaa)
      {}

      node_iterator& operator++ ();
      node_iterator operator++ (int);
      NumericArrayType* operator* ();
      bool operator== (const node_iterator&);
      bool operator!= (const node_iterator&);
      bool end_p ();
      
      node_iterator& pp (int active_idx);
      node_iterator& mm (int active_idx);
      NumericType& data ();
      void refine (int rel_level = 1);
      void coarse (int rel_level = 1);
      int level ();
      IndexType& position ();
      int position (int idx);
      AdaptiveArray* get_aa ();
      ArrayType* base ();
      Node*& node ();

      level_iterator level_iter (int); 
    };

    friend class node_iterator;

  public:
    AdaptiveArray (IndexType& size,  NumericType init, int iRefineLevel = 4); 

    NumericType& operator() (IndexType& pos);
    NumericType& operator() (IndexType& npos, IndexType& vpos);

    node_iterator begin ();
    node_iterator end ();
    node_iterator at (const IndexType& pos, bool finest = true);
    
    IndexType& dims ();

    NumericType& data (IndexType& pos);

    inline int refine_level ();

  public:
    static IndexType rel_position (IndexType& finest_pos);

  protected:
    void refine (Node*& node, NumericType val, int rel_level = 1);
    void coarse (Node*& node, NumericType& val, int rel_level = 1);
  };


  /// Implementation

  template<typename ref_coars_constr>
  AdaptiveArray<ref_coars_constr>::AdaptiveArray (IndexType& isize, 
						  typename ref_coars_constr::NumericType init,
						  int iRefineLevel)
    : RefineLevel (iRefineLevel)
  {
    for (int i = 0; i < ref_coars_constr::Dim; i++)
      {
	size [i] = (isize [i] >> RefineLevel);
	if (isize [i] & ((1 << RefineLevel) - 1))
	  size [i] ++;
      }

    node_array.reference (ArrayType (size));
    data_array.reference (NumericArrayType (size));
    data_array = init;
  }

  template<typename ref_coars_constr>
  typename ref_coars_constr::NumericType& 
  AdaptiveArray<ref_coars_constr>::operator () (IndexType& pos)
  {
    IndexType p, p1;
 
#if 0
    for (int i = 0; i < dim; i ++)
      p [i] = pos [i] >> max_ref_lev;
#else
    p = pos >> RefineLevel;
#endif
 
    Node* n = node_array (p);

    if (!n)
      return data_array (p);

    int sub = RefineLevel - n->level;
    
#if 0
    for (int i = 0; i < ref_coars_constr::Dim; i++)
      p1 [i] = (pos [i] % (1 << RefineLevel)) >> sub;
#else
    p1 = (pos % (1 << RefineLevel)) >> sub;
#endif
 
    return n->data (p1);
  }
  
  template<typename ref_coars_constr>
  typename ref_coars_constr::NumericType& 
  AdaptiveArray<ref_coars_constr>::operator () (IndexType& npos, IndexType& vpos)
  {
    Node* n = node_array (npos);

    if (n == NULL)
      return data_array (npos);

    return n->data (vpos);
  }

  template<typename ref_coars_constr>
  void AdaptiveArray<ref_coars_constr>::refine (AdaptiveArray::Node*& node, 
						typename ref_coars_constr::NumericType val,
						int rel_level)
  {
    if (!node)
      node = new typename AdaptiveArray<ref_coars_constr>::Node;

    if (node->level == RefineLevel)
      return;

    if ((node->level + rel_level) > RefineLevel)
      rel_level = RefineLevel - node->level;

    int size = 1 << node->level;
    NumericArrayType new_data (IndexType (size << rel_level));

    if (size == 1)
      ref_coars_constr::init_refine (new_data, val);
    else
      {
	for (typename NumericArrayType::iterator i = node->data.begin (); i != node->data.end (); ++i)
	  ref_coars_constr::init_refine (new_data, i, rel_level);
      }

    node->data.reference (new_data);
    node->level += rel_level;
  }

  template<typename ref_coars_constr>
  void AdaptiveArray<ref_coars_constr>::coarse (AdaptiveArray::Node*& node, 
						typename ref_coars_constr::NumericType& val,
						int rel_level)
  {
    if (!node->level)
      return;

    if ((node->level - rel_level) <= 0)
      {
	ref_coars_constr::init_coarse (val, node->data);
	delete node;
	node = NULL;
	return;
      }

    int size = 1 << node->level;
    NumericArrayType new_data (IndexType (size >> rel_level));

    for (typename NumericArrayType::iterator i = node->data.begin (); i != node->data.end (); ++i)
      ref_coars_constr::init_coarse (new_data, i, rel_level);

    node->data.reference (new_data);
    node->level -= rel_level;
  }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::node_iterator
  AdaptiveArray<ref_coars_constr>::begin ()
  { return typename AdaptiveArray<ref_coars_constr >::node_iterator (this); }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr >::node_iterator
  AdaptiveArray<ref_coars_constr>::end ()
  { return typename AdaptiveArray<ref_coars_constr >::node_iterator (this, size); }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr >::node_iterator
  AdaptiveArray<ref_coars_constr>::at (const IndexType& pos, bool finest)
  { 
    if (finest)
      {
	blitz::TinyVector<int, ref_coars_constr::Dim> tmp (pos >> RefineLevel);
	return typename AdaptiveArray::node_iterator (this, tmp);
      }
    else
      return typename AdaptiveArray::node_iterator (this, pos);
  }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr >::IndexType&
  AdaptiveArray<ref_coars_constr>::dims ()
  { return size; }

  template<typename ref_coars_constr>
  typename ref_coars_constr::NumericType& 
  AdaptiveArray<ref_coars_constr>::data (IndexType& pos)
  { return data_array (pos); }

  template<typename ref_coars_constr>
  inline int
  AdaptiveArray<ref_coars_constr>::refine_level ()
  { return RefineLevel; }


  /// Iterators

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::node_iterator& 
  AdaptiveArray<ref_coars_constr>::node_iterator::operator++ ()
  {
    int i = ref_coars_constr::Dim - 1;

    pos [i] ++;
    while (pos [i] == aa->size [i])
      {
	pos [i] = 0;
	i --;
	if (i < 0)
	  {
	    pos = aa->size;
	    break;
	  }
	pos [i] ++;
      }

    return *this;
  }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::node_iterator 
  AdaptiveArray<ref_coars_constr>::node_iterator::operator++ (int)
  {
    typename AdaptiveArray<ref_coars_constr>::node_iterator tmp = *this;
    ++(*this);
    return tmp;
  }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::node_iterator& 
  AdaptiveArray<ref_coars_constr>::node_iterator::pp (int active_idx)
  {
    pos [active_idx] ++;
    if (pos [active_idx] >= aa->size [active_idx])
      pos = aa->size;
    
    return *this;
  }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::node_iterator& 
  AdaptiveArray<ref_coars_constr>::node_iterator::mm (int active_idx)
  {
    pos [active_idx] --;
    if (pos [active_idx] < 0)
      pos = aa->size;
    
    return *this;
  }

  template<typename ref_coars_constr>
  void
  AdaptiveArray<ref_coars_constr>::node_iterator::refine (int rel_level)
  { aa->refine (aa->node_array (pos), aa->data_array (pos), rel_level); }

  template<typename ref_coars_constr>
  void
  AdaptiveArray<ref_coars_constr>::node_iterator::coarse (int rel_level)
  { aa->coarse (aa->node_array (pos), aa->data_array (pos), rel_level); }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::NumericArrayType*
  AdaptiveArray<ref_coars_constr>::node_iterator::operator* ()
  { 
    if (aa->node_array (pos))
      return &(aa->node_array (pos)->data); 
    return NULL;
  }

  template<typename ref_coars_constr>
  bool
  AdaptiveArray<ref_coars_constr>::node_iterator::operator== 
  (const AdaptiveArray<ref_coars_constr>::node_iterator& it)
  {
    if (aa != it.aa)
      return false;

    return eq_index<int, ref_coars_constr::Dim> (pos, it.pos);
  }

  template<typename ref_coars_constr>
  bool
  AdaptiveArray<ref_coars_constr>::node_iterator::operator!= 
  (const AdaptiveArray<ref_coars_constr>::node_iterator& it)
  {
    if (aa != it.aa)
      return true;

    return neq_index<int, ref_coars_constr::Dim> (pos, it.pos);
  }

  template<typename ref_coars_constr>
  bool
  AdaptiveArray<ref_coars_constr>::node_iterator::end_p ()
  { return neq_index<int, ref_coars_constr::Dim> (pos, aa->dims ()); }

  template<typename ref_coars_constr>
  typename ref_coars_constr::NumericType&
  AdaptiveArray<ref_coars_constr>::node_iterator::data ()
  { return aa->data_array (pos); }

  template<typename ref_coars_constr>
  int
  AdaptiveArray<ref_coars_constr>::node_iterator::level ()
  {
    Node* n = aa->node_array (pos);

    if (!n)
      return 0;

    return n->level; 
  }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::IndexType&
  AdaptiveArray<ref_coars_constr>::node_iterator::position ()
  { return pos; }

  template<typename ref_coars_constr>
  int
  AdaptiveArray<ref_coars_constr>::node_iterator::position (int idx)
  { return pos [idx]; }

  template<typename ref_coars_constr>
  AdaptiveArray<ref_coars_constr> *
  AdaptiveArray<ref_coars_constr>::node_iterator::get_aa ()
  { return aa; }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::ArrayType*
  AdaptiveArray<ref_coars_constr>::node_iterator::base ()
  { return &aa->node_array; }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::Node*&
  AdaptiveArray<ref_coars_constr>::node_iterator::node ()
  { return  (aa->node_array (pos)); }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::level_iterator
  AdaptiveArray<ref_coars_constr>::node_iterator::level_iter (int level_up)
  { return typename AdaptiveArray<ref_coars_constr>::level_iterator (*this, level_up); }

  ///

  template<typename ref_coars_constr>
  AdaptiveArray<ref_coars_constr>::level_iterator::level_iterator (AdaptiveArray<ref_coars_constr>::node_iterator& iit, int ilevel_up)
    : pos (blitz::TinyVector<int, ref_coars_constr::Dim> (0), blitz::TinyVector<int, ref_coars_constr::Dim> ((*iit)->ubound () + 1)),
      it (iit), level_up (ilevel_up)
  {}

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::level_iterator&
  AdaptiveArray<ref_coars_constr>::level_iterator::operator++ ()
  {
    pos += (1 << level_up);
    
    return (*this);
  }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::NumericArrayType
  AdaptiveArray<ref_coars_constr>::level_iterator::operator* ()
  {
    blitz::RectDomain<ref_coars_constr::Dim> d (*pos, *pos + (1 << level_up) - 1);
    
    return (*(*it)) (d);
  }

  template<typename ref_coars_constr>
  bool
  AdaptiveArray<ref_coars_constr>::level_iterator::end_p ()
  { return pos.end_p (); }

  /// 

  template<typename ref_coars_constr>
  AdaptiveArray<ref_coars_constr>::cell_index::cell_index 
  (typename AdaptiveArray<ref_coars_constr>::IndexType node,
   typename AdaptiveArray<ref_coars_constr>::IndexType val)
    : node_idx (node), val_idx (val)
  {}

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::IndexType&
  AdaptiveArray<ref_coars_constr>::cell_index::node ()
  { return node_idx; }

  template<typename ref_coars_constr>
  typename AdaptiveArray<ref_coars_constr>::IndexType&
  AdaptiveArray<ref_coars_constr>::cell_index::val ()
  { return val_idx; }

  /// Constraints

  template<typename num_type, int dim>
  class empty_constr
  {};

  template<typename num_type, int dim>
  class aa_equal_distr_constr  : public empty_constr<num_type, dim>
  {
  public:
    typedef num_type NumericType;

    typedef aa_equal_distr_constr<num_type, dim> Self;
    typedef AdaptiveArray<Self> ArrayType;
    typedef blitz::TinyVector<int, dim> IndexType;
    typedef blitz::Array<NumericType, dim> NumericArrayType;

    enum { Dim = dim };

  public:
    static void init_refine (NumericArrayType& new_array, num_type& val)
    { new_array = val; }

    static void init_refine (NumericArrayType& new_array, 
			     typename NumericArrayType::iterator& it,
			     int rel_level = 1)
    {
      IndexType pos (it.position () << rel_level);
      blitz::RectDomain<dim> r (pos, pos + (1 << rel_level) - 1);
      num_type val = *it;

      new_array (r) = val;
    }

    static void init_coarse (num_type& val, NumericArrayType& new_array)
    { val =  new_array (IndexType (0)); }

    static void init_coarse (NumericArrayType& new_array, 
			     typename NumericArrayType::iterator& it,
			     int rel_level = 1)
    { new_array (IndexType (it.position () >> rel_level)) = *it; }

    static bool coarse_p (NumericArrayType& new_array)
    {
      num_type val = new_array (IndexType (0));

      for (typename NumericArrayType::iterator it = new_array.begin (); it != new_array.end (); ++it)
	if (val != *it)
	  return false;

      return true;
    }
  };

}; // namespace dbcc

#endif
